# Copyright (C) 2007 Adriano Monteiro Marques <py.adriano@gmail.com>
#
# Author: Maxim Gavrilov <lovelymax@gmail.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

import os, os.path
import re
import sys

def check_access(path, permission):
    return os.path.exists(path) and os.access(path, permission)

def get_file_list(path):
        result = []
        try:
            for filename in os.listdir(path):
                fullpath = os.path.join(path, filename)
                if os.path.isfile(fullpath) and check_access(fullpath, os.R_OK):
                    result.append(fullpath)
        except OSError:
            pass
        return result

# aux nmap_fetchfile functions
class NmapFetch(object):
    # Singletone for optimization
    _instance = None
    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(NmapFetch, cls).__new__(cls, *args, **kwargs)
        return cls._instance
    
    def __init__(self):
        self.dirs = self.__fetchdirs()

    def fetchdirs(self):
        return self.dirs

    def fetchfile(self, filename):
        for path in self.fetchdirs():
            fullpath = os.path.join(path, filename)
            if check_access(fullpath, os.R_OK):
                return fullpath
        return None

    def get_file_list(self):
        result = []
        for path in self.fetchdirs():
            result.extend(get_file_list(path))
        return result

    def nmap_path(self, path):
        fullpath = os.path.abspath(path)
        for p in self.fetchdirs():
            if fullpath.startswith(p):
                return fullpath[len(p)+1:] # XXX: check +1 (removing last slash) on Windows
        return fullpath

    def __fetchdirs(self):
        # standart Nmap searching directories (see nmap.cc:nmap_fetchfile function)
        def varpath():
            return os.path.expandvars("${NMAPDIR}")
        def uidpath():
            return os.path.join(pwd.getpwuid(os.getuid()).pw_dir, ".nmap")
        def euidpath():
            return os.path.join(pwd.getpwuid(os.geteuid()).pw_dir, ".nmap")
        def userpath():
            return os.path.expanduser("~")
        def datadirpath_win():
            return "c:\\nmap"
        def datadirpath2_win():
            return "c:\\Program Files\\nmap"
        def datadirpath():
            return "/usr/share/nmap/"
        def datadirpath2():
            return "/usr/local/share/nmap/"
        def currentpath():
            return "."

        if sys.platform != 'win32':
            import pwd
            checklist = [varpath, uidpath, euidpath, datadirpath, datadirpath2, currentpath]
        else:
            checklist = [varpath, userpath, datadirpath_win, datadirpath2_win, currentpath]

        paths = [os.path.abspath(f()) for f in checklist]
        # XXX: not stable
        return list(set(paths))

class NmapFetchScripts(NmapFetch):
    def __init__(self):
        NmapFetch.__init__(self)
        self.dirs = [os.path.join(d, "scripts") for d in self.dirs]

    def get_file_list(self):
        return [f for f in NmapFetch.get_file_list(self) if f.endswith(".nse")]

class ScriptParseException(Exception):
    pass

class Script(object):
    def __init__(self, path):
        f = file(path, 'r')
        self.data = f.read()
        f.close()

        self.path = NmapFetchScripts().nmap_path(path)
        self.name, self.type = os.path.splitext(os.path.basename(path))
        self.id = self._get_attr('id') # XXX: raise exception if no id in script
        self.desc = self._get_attr('description', "No descriptions")
        self.author = self._get_attr("author", "unnamed")

        self.categories = set(["untagged"])
        self.categories.update(set(self._get_attr_list("categories", set())))
        # HTTPAuth used tags as local variable
        #self.categories.update(set(self._get_attr_list("tags", set())))
        if len(self.categories) > 1:
            self.categories.remove("untagged")

    def _normilize(self, string):
        string = string.replace('\n', ' ')
        string = string.replace('\\', ' ')
        string = string.replace('   ', ' ')
        string = string.replace('   ', '\n')
        return " ".join(string.split(" ")).strip()
        # more script-related breaks
        #return "\n".join([" ".join(s.split(" ")).strip()
        #                  for s in string.replace('\\', '\n').split("\n")])
        
    def _get_attr_list(self, attr, default = None):
        r = re.findall(attr + '\s*=\s*{([^\}]*)}', self.data)
        if not r:
            if default is not None:
                return default
            raise ScriptParseException("Can't parse attr %s in file %s" % (attr, self.path))
        return [self._normilize(tag) for tag in re.findall('"([^\"]+)"', r[0])]
        
    def _get_attr(self, attr, default = None):
        r = re.findall(attr + '\s*=\s*"([^\"]+)"', self.data)
        if not r:
            if default is not None:
                return default
            raise ScriptParseException("Can't parse attr %s in file %s" % (attr, self.path))
        return self._normilize(r[0])

    # for set-element support
    def __eq__(self, other):
        return self.path.__eq__(other.path)

    def __hash__(self):
        return self.path.__hash__()

class ScriptSelection(object):
    def __init__(self, script_manager, selected = ""):
        self.script_manager = script_manager
        self.d = dict([(s, False) for s in self.script_manager])
        self.set_selected(selected)

    def set_selected(self, selected):
        for filename in selected.split(";"):
            script = self.script_manager.find_by_path(filename)
            if self.d.has_key(script):
                self.d[script] = True

    def get_selected(self):
        return ";".join([script.path for script in self.d.keys() if self.d[script]])

    def is_selected(self, script):
        return self.d.get(script, False)

    def select(self, script):
        self.d[script] = True

    def unselect(self, script):
        self.d[script] = False

# Script Sources
class ScriptSource(object):
    def __init__(self):
        self.scripts = None
        
    def update_scripts(self):
        if not self.scripts:
            self.scripts = self.get_scripts_impl()
        return self.scripts
    
    def reload_scripts(self):
        self.scripts = self.get_scripts_impl()
        return self.scripts
        
class ScriptFile(ScriptSource):
    def __init__(self, filename):
        ScriptSource.__init__(self)
        self.filename = filename

    def get_scripts_impl(self):
        res = []
        try:
            res = [Script(self.filename)]
        except ScriptParseException:
            pass
        return res
 
class ScriptDir(ScriptSource):
    def __init__(self, dirname):
        ScriptSource.__init__(self)
        self.dirname = dirname

    def get_scripts_impl(self):
        res = []
        for name in get_file_list(self.dirname):
            if name.endswith('.nse'):
                fullname = os.path.join(self.dirname, name)
                try:
                    res.append(Script(fullname))
                except ScriptParseException:
                    pass
        return res
    
# Script Manager                    
class ScriptManager(set):
    def __init__(self):
        self.sources = []
        self.categories = set([])
        for dirname in NmapFetchScripts().fetchdirs():
            self.sources.append(ScriptDir(dirname))
        self.selection = ScriptSelection(self)
        self.update_scripts()

    def add_file(self, filename):
        self.sources.append(ScriptFile(filename))
        self.update_scripts()

    def add_dir(self, dirname):
        src = ScriptDir(dirname)
        self.sources.append(src)
        self.update_scripts()
        return len(src.get_scripts()) > 0

    def find_by_path(self, filename):
        for s in self:
            if s.path == filename:
                return s
        return None

    def update_categories(self):
        self.categories = set([])
        for s in self:
            self.categories.update(s.categories)

    def update_scripts(self):
        self.clear()
        for src in self.sources:
            self.update(src.update_scripts())
        self.update_categories()

    def reload_scripts(self):
        self.clear()
        for src in self.sources:
            self.update(src.reload_scripts())
        self.update_categories()

    def get_selection(self):
        return self.selection

if __name__ == "__main__":
    manager = ScriptManager()
    for script in manager:
        print script.name
    
